@kinianlo Thanks for this, at first glance looks really good! @nikhilkhatri will have a detailed look shortly, but in the meantime, it would be great if you could post a notebook that shows how it works in practice?

If you can also rebase on the current branch it would be great.

Sure here is an example how the new model can be used:

from lambeq.backend.grammar import Cup, Id, Word
from lambeq import AtomicType, IQPAnsatz, PytorchQuantumModel
import torch
torch.manual_seed(42)

N = AtomicType.NOUN
S = AtomicType.SENTENCE

ansatz = IQPAnsatz({N: 1, S: 1}, n_layers=1)
diagrams = [ansatz((Word("Alice", N) @ Word("runs", N >> S) >> Cup(N, N.r) @ Id(S))),
            ansatz((Word("Bob", N) @ Word("jumps", N >> S) >> Cup(N, N.r) @ Id(S)))]
labels = torch.tensor([0, 1])

model = PytorchQuantumModel.from_diagrams(diagrams)
model.initialise_weights()

optimizer = torch.optim.Adam(model.parameters(), lr=1e-2)
loss_fn = torch.nn.CrossEntropyLoss()
    
for epoch in range(10):
    optimizer.zero_grad()
    pred = model(diagrams)
    loss = loss_fn(pred, labels)
    loss.backward()
    optimizer.step()
    print(f'Epoch {epoch}, Loss: {loss.item()}')

which prints

Epoch 0, Loss: 0.5054250025233994
Epoch 1, Loss: 0.4690612002974437
Epoch 2, Loss: 0.44169619906825297
Epoch 3, Loss: 0.42269240175382544
Epoch 4, Loss: 0.40758353898053123
Epoch 5, Loss: 0.39317936341381804
Epoch 6, Loss: 0.3788896391036708
Epoch 7, Loss: 0.364901423245132
Epoch 8, Loss: 0.3516455131919579
Epoch 9, Loss: 0.33967176736231863

Another example with PytorchTrainer:

from lambeq.backend.grammar import Cup, Id, Word
from lambeq import AtomicType, IQPAnsatz, PytorchQuantumModel
from lambeq import PytorchTrainer, Dataset
import torch
torch.manual_seed(42)

N = AtomicType.NOUN
S = AtomicType.SENTENCE

ansatz = IQPAnsatz({N: 1, S: 1}, n_layers=1)
diagrams = [ansatz((Word("Alice", N) @ Word("runs", N >> S) >> Cup(N, N.r) @ Id(S))),
            ansatz((Word("Bob", N) @ Word("jumps", N >> S) >> Cup(N, N.r) @ Id(S)))]
labels = [[1, 0], [0, 1]]

model = PytorchQuantumModel.from_diagrams(diagrams)
model.initialise_weights()
trainer = PytorchTrainer(model, 
                         loss_function=torch.nn.CrossEntropyLoss(),
                         optimizer=torch.optim.Adam,
                         epochs=10,
                         learning_rate=1e-2)
dataset = Dataset(diagrams, labels)

trainer.fit(dataset)

which prints:

Epoch 1:   train/loss: 0.8319   valid/loss: -----   train/time: 0.02s   valid/time: -----
Epoch 2:   train/loss: 0.7689   valid/loss: -----   train/time: 0.01s   valid/time: -----
Epoch 3:   train/loss: 0.7070   valid/loss: -----   train/time: 0.01s   valid/time: -----
Epoch 4:   train/loss: 0.6491   valid/loss: -----   train/time: 0.01s   valid/time: -----
Epoch 5:   train/loss: 0.5979   valid/loss: -----   train/time: 0.01s   valid/time: -----
Epoch 6:   train/loss: 0.5553   valid/loss: -----   train/time: 0.01s   valid/time: -----
Epoch 7:   train/loss: 0.5215   valid/loss: -----   train/time: 0.01s   valid/time: -----
Epoch 8:   train/loss: 0.4960   valid/loss: -----   train/time: 0.01s   valid/time: -----
Epoch 9:   train/loss: 0.4780   valid/loss: -----   train/time: 0.01s   valid/time: -----
Epoch 10:  train/loss: 0.4666   valid/loss: -----   train/time: 0.01s   valid/time: -----

Training completed!
train/time: 0.08s   train/time_per_epoch: 0.01s   train/time_per_step: 0.01s   valid/time: None   valid/time_per_eval: None

If you can also rebase on the current branch it would be great.

Done!

The current implementation of PytorchQuantumModel is somewhat hacky. It inherits from both PytorchModel and QuantumModel:

1. PytorchModel provides access to _load_checkpoint, _load_checkpoint and _reinitialise_modules,
2. QuantumModel provides access to _normalise_vector and _fast_subs.

This approach works but definitely not the best design and it introduces all sorts of type check errors. Perhaps @nikhilkhatri would have a better idea on how to structure the inheritance.

Thanks a lot for this PR @kinianlo, it looks largely very good!
I have a few broad comments:

1. Could you please add a test for training with mixed states?
   • See an example where such a circuit is created here: https://github.com/CQCL/lambeq/blob/main/tests/training/test_numpy_model.py#L49
2. Re. the multiple inheritance causing type checking failures:
   • For now, I'd suggest only inheriting from QuantumModel, and duplicating the functions from PytorchModel
   • Those functions are pretty lightweight, and maybe they can be abstracted away in a future PR.

Thanks for the feedback @nikhilkhatri. I have added the test for mixed states and removed PytorchModel for the inheritance.

However the type problems still stand. Mostly because QuantumModel was made for numpy arrays in mind and typed accordingly. There are so many different parts of the codebase which expect to see only numpy arrays from QuantumModel. One example is SPSAOptimizer. It seems like a lot of work has to be done in order to make QuantumModel compatible with pytorch.

Ofcourse one way out would be to not inherit QuantumModel but that doesn't feel right.

Hi @kinianlo, great work on this PR, thank you!

Could you rebase your branch with the latest main branch to fix the issues with the GH workflow? Thank you very much. I'm also reviewing your work about the inheritance issues you mentioned above.

@neiljdo Thanks! I have rebased the branch.

@kinianlo sorry for the trouble but can you rebase again with the latest main branch? I just merged a PR that should solve the typing issues related to subclassing from the QuantumModel class.

EDIT: I've rebased and updated your branch.